from fastapi import FastAPI, Body
from pydantic import BaseModel
from typing import Dict
import base64
from PIL import Image
import io
from SensitiveImgDetect import Detect

app = FastAPI()
detector = Detect(device='cpu')  # 初始化模型，根据你的环境选择 'cuda' 或 'cpu'

class ImageRequest(BaseModel):
    base64_image: str



@app.post("/detect_image/")
async def detect_image(request: ImageRequest):
    """
    接收 Base64 编码的图片，使用 SensitiveImgDetect 模型检测其类别和概率。
    """
    try:
        # 解码 Base64 字符串
        image_bytes = base64.b64decode(request.base64_image)
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        # 检测图片类别和概率
        predicted_label = detector.detect_single_type(image)
        probabilities = detector.detect_single_prob(image)

        return {
            "status_code": 200,
            "msg": "检测成功",
            "data": {
                "predicted_class": predicted_label,
                "probabilities": probabilities
            }
        }

    except Exception as e:
        # return {"error": f"图像处理或检测失败: {e}"}
        return {
            "status_code": 500,
            "msg": f"图像处理或检测失败: {e}",
            "data": None
        }

# uvicorn main:app --reload

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8048)
    
